from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import JsonOutputParser
from pydantic import BaseModel, Field
from typing import List


class Plan(BaseModel):
    """
    制定的计划
    """

    task_list: List[str] = Field(description="子任务列表")


_planner_system_template = """
“您是一位卓越的任务规划专家，您的职责是为用户提出的查询目标,用最简方式制定一个全面而细致的解答方案。”
输出:
{output_format}

"""
_planner_human_template = """
查询目标:
{query}

"""


class Planner:
    def __init__(self, llm):
        _prompt = ChatPromptTemplate.from_messages(
            [
                ("system", _planner_system_template),
                ("human", _planner_human_template),
            ]
        )

        _parser = JsonOutputParser(pydantic_object=Plan)
        _prompt = _prompt.partial(output_format=_parser.get_format_instructions())

        self._chain = _prompt | llm | _parser

    def __call__(self, state):
        return self._chain.invoke(state)


if __name__ == "__main__":
    from langchain_openai import ChatOpenAI

    _llm = ChatOpenAI(
        base_url="http://192.168.10.11:60026/v1",
        model="qwen2.5:7b",
        api_key="ollama",
        temperature=0.4,
    )
    planner = Planner(_llm)
    _rt = planner({"query":"2024年法国巴黎奥运会女子10米跳台冠军的父母是谁？"})
    print(_rt)
